package com.jboltai.capability.model;


import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSON;
import com.jboltai.capability.Templates;
import com.jboltai.capability.message.AIIntention;
import com.jboltai.capability.message.AIMessage;
import com.jboltai.capability.vdb.doc.VDBDoc;
import com.jboltai.config.JBoltAIConfig;
import com.jboltai.event.ai.aichat.AIChatEvent;
import com.jboltai.event.ai.aifunccall.AIFuncCallEvent;
import com.jboltai.event.ai.aiintention.AIIntentionEvent;
import com.jboltai.event.ai.aiknowledge.AIKnowledgeEvent;
import com.jboltai.event.ai.aitext.text2json.Text2JsonEvent;
import com.jboltai.event.ai.aitext.text2sql.DBType;
import com.jboltai.event.ai.aitext.text2sql.SqlType;
import com.jboltai.event.ai.aitext.text2sql.Text2SqlEvent;
import com.jboltai.event.ai.rewrite.AIRewritePromptEvent;
import com.jboltai.event.chain.Chain;
import com.jboltai.resource.ai.AIModel;
import com.jboltai.resource.ai.AIResource;
import com.jboltai.resource.ai.AIResourceCenter;
import com.jboltai.scheduler.EventBus;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.math.BigDecimal;
import java.util.*;

/**
 * AI入口
 */
public class JBoltAI {

    private static final Logger LOGGER = LoggerFactory.getLogger(JBoltAIConfig.LOGGER_NAME);

    /**
     * 注册一个ai模型资源，返回一个资源对象，对象id随机生成
     *
     * @param model
     * @param apiKey
     * @return
     */
    public static AIResource registerResource(AIModel model, String apiKey) {
        AIResource resource = new AIResource(model, apiKey);
        AIResourceCenter.register(resource);
        return resource;
    }

    /**
     * 注册一个ai模型资源，返回一个资源对象，对象id使用指定的id
     *
     * @param model
     * @param id
     * @param apiKey
     * @return
     */
    public static AIResource registerResource(AIModel model, String id, String apiKey) {
        AIResource resource = new AIResource(model, id, apiKey);
        AIResourceCenter.register(resource);
        return resource;
    }

    /**
     * 注册一个ai模型资源，返回一个资源对象，对象id使用指定的id
     *
     * @param model
     * @param id
     * @param apiKey
     * @return
     */
    public static AIResource registerResource(AIModel model, String id, String apiKey, String secretKey) {
        AIResource resource = new AIResource(model, id, apiKey, secretKey);
        AIResourceCenter.register(resource);
        return resource;
    }

    public static AIResource registerResource(AIModel model, String id, String appId, String apiKey, String secretKey) {
        AIResource resource = new AIResource(model, id, appId, apiKey, secretKey);
        AIResourceCenter.register(resource);
        return resource;
    }

    /**
     * 根据id，取消注册某个ai资源
     *
     * @param id
     */
    public static void unregisterResource(String id) {
        AIResource resource = AIResourceCenter.getResourceById(id);
        if (resource != null) {
            AIResourceCenter.unregister(resource);
        }
    }


    public static Chain chain() {
        return new Chain();
    }


    /**
     * 构建一个ai聊天事件
     *
     * @return
     */
    public static AIChatEvent chat() {
        return new AIChatEvent();
    }


    /**
     * 使用指定的模型构建一个ai聊天事件
     *
     * @param model
     * @return
     */
    public static AIChatEvent chat(AIModel model) {
        return new AIChatEvent(model);
    }


    /**
     * 构建一个function call事件
     *
     * @return
     */
    public static AIFuncCallEvent funcCall() {
        return new AIFuncCallEvent();
    }


    /**
     * 使用指定的模型构建一个function call聊天事件
     *
     * @param model
     * @return
     */
    public static AIFuncCallEvent funcCall(AIModel model) {
        return new AIFuncCallEvent(model);
    }


    /**
     * 将functionCall的结果和问题结合后给出回答
     *
     * @param question           问题
     * @param funcDesc           业务描述
     * @param funcCallResult     functionCall结果
     * @param returnRequirements 返回数据要求
     * @return
     */
    public static AIChatEvent chatWithFuncCallResult(String question, String funcDesc, Object funcCallResult, String returnRequirements) {
        String result = funcCallResult == null ? "" : (funcCallResult instanceof String || funcCallResult instanceof Number) ? funcCallResult.toString() : JSON.toJSONString(funcCallResult);
        String processMessage = StrUtil.format(Templates.CHAT_WITH_FUNCTION_CALL,funcDesc,result,question,StrUtil.blankToDefault(returnRequirements, StrUtil.EMPTY));
        LOGGER.debug("chatWithFuncCallResult:{}", processMessage);
        AIChatEvent event = new AIChatEvent();
        event.prompt(processMessage);
        return event;
    }

    /**
     * 将functionCall的结果和问题结合后给出回答
     *
     * @param question       问题
     * @param funcDesc       业务描述
     * @param funcCallResult functionCall结果
     * @return
     */
    public static AIChatEvent chatWithFuncCallResult(String question, String funcDesc, Object funcCallResult) {
        return chatWithFuncCallResult(question, funcDesc, funcCallResult, null);
    }


    /**
     * 运用知识库给出回答
     *
     * @param question
     * @param background    源自知识库的背景知识
     * @param knowledgeOnly 是否仅基于给定的知识库进行回答
     * @return
     */
    public static AIKnowledgeEvent chatWithKnowledge(String question, List<String> background, boolean knowledgeOnly) {
        AIKnowledgeEvent event = new AIKnowledgeEvent();
        event.prompt(question).setBackground(background).setKnowledgeOnly(knowledgeOnly);
        return event;
    }

    /**
     * 运用知识库给出回答
     *
     * @param promptTpl  自定义知识库模版，模板中支持插入三个变量：{background} 背景知识，{question} 问题，{nomatch} 无匹配时的回答
     * @param question
     * @param background 源自知识库的背景知识
     * @return
     */
    public static AIKnowledgeEvent chatWithKnowledge(String promptTpl, String question, List<String> background, String nomatchReply) {
        AIKnowledgeEvent event = new AIKnowledgeEvent();
        event.prompt(question).setBackground(background).setPromptTpl(promptTpl).setNoMatchReply(nomatchReply);
        return event;
    }

    /**
     * 运用知识库给出回答
     *
     * @param
     * @param question
     * @param background   源自知识库的背景知识
     * @param nomatchReply 无匹配时的回答
     * @return
     */
    public static AIKnowledgeEvent chatWithKnowledge(String question, List<String> background, String nomatchReply) {
        AIKnowledgeEvent event = new AIKnowledgeEvent();
        event.prompt(question).setBackground(background).setNoMatchReply(nomatchReply).setNoMatchReply(nomatchReply);
        return event;
    }

    /**
     * 根据用户提问，进行意图识别
     * @param prompt
     * @param intentions
     * @return 返回匹配到的意图的key，如果未匹配到任何给定的意图，将返回null
     */
    public static AIIntentionEvent recognizeIntention(String prompt, AIIntention... intentions) {
        AIIntentionEvent event = new AIIntentionEvent();
        event.prompt(prompt);
        event.bindIntentionOptions(intentions );
        return event;
    }


    /**
     * 根据用户问题生成sql语句
     *
     * @return
     */
    public static Text2SqlEvent text2Sql() {
        return new Text2SqlEvent();
    }


    /**
     * 根据用户问题生成sql语句
     *
     * @param dbType 数据库类型
     * @return
     */
    public static Text2SqlEvent text2Sql(DBType dbType) {
        return new Text2SqlEvent().dbType(dbType);
    }

    /**
     * 根据用户问题生成sql语句
     *
     * @param dbType     数据库类型
     * @param background 背景知识
     * @return
     */
    public static Text2SqlEvent text2Sql(DBType dbType, String background, SqlType... limitSqlType) {
        return new Text2SqlEvent().dbType(dbType).background(background).limitSqlType(limitSqlType);
    }


    /**
     * 根据用户内容生成json结构数据
     *
     * @return
     */
    public static Text2JsonEvent text2Json() {
        return new Text2JsonEvent();
    }


    /**
     * 据用户内容生成json结构数据
     *
     * @param background 背景知识
     * @param jsonDesc   json结构 说明
     * @return
     */
    public static Text2JsonEvent text2Json(String background, String jsonDesc) {
        if (StrUtil.isBlank(jsonDesc)) {
            throw new IllegalArgumentException("jsonDesc不能为空");
        }
        return new Text2JsonEvent().background(background).jsonDesc(jsonDesc);
    }

    /**
     * 根据id，取消某个事件
     *
     * @param eventId
     */
    public static void cancel(String eventId) {
        EventBus.cancel(eventId);
    }


    /**
     * 使用rrf算法进行vdb结果集的重排, 默认k=60
     * @param docCollections
     */
    public static List<VDBDoc> rerankByRRF(List<VDBDoc>... docCollections) {
        return rerankByRRF(60, docCollections);
    }

    /**
     * 使用rrf算法进行vdb结果集的重排
     * @param k rrf算法的参数k，用于平衡各结果集的权重
     * @param docCollections
     */
    public static List<VDBDoc> rerankByRRF(int k, List<VDBDoc>... docCollections) {
        Map<String, BigDecimal> scores = new HashMap<>();
        List<VDBDoc> list = new ArrayList<>();

        for (List<VDBDoc> docCollection : docCollections) {
            for (int i = 0; i < docCollection.size(); i++) {
                VDBDoc doc = docCollection.get(i);
                if (scores.containsKey(doc.getId())) {
                    VDBDoc oriDoc = list.stream().filter(item -> item.getId().equals(doc.getId())).findFirst().get();
                    if (oriDoc.getScore() != null && doc.getScore() != null && doc.getScore() > oriDoc.getScore()) {
                        oriDoc.setScore(doc.getScore());
                    } else if (oriDoc.getScore() == null && doc.getScore() != null) {
                        oriDoc.setScore(doc.getScore());
                    }


                    scores.put(doc.getId(), scores.get(doc.getId()).add(new BigDecimal(1.0 / (k + i + 1))));
                } else {
                    scores.put(doc.getId(), new BigDecimal(1.0 / (k + i + 1)));
                    list.add(doc);
                }
            }
        }
        list.sort((a, b) -> {
            return scores.get(b.getId()).compareTo(scores.get(a.getId())) ;
        });

        return list;

    }

    /**
     * 构建一个重写Prompt的事件
     *
     * @return
     */
    public static AIRewritePromptEvent rewritePrompt() {
        return new AIRewritePromptEvent();
    }

    /**
     * 构建一个重写Prompt的事件
     * @param prompt
     * @return
     */
    public static AIRewritePromptEvent rewritePrompt(String prompt, List<AIMessage> contextMsgs) {
        return new AIRewritePromptEvent().prompt(prompt).contextMessages(contextMsgs);
    }


}
